package com.wzf.greattruth.zookeeper.queue;

import org.apache.commons.lang3.StringUtils;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.ZooKeeper;

import com.wzf.greattruth.zookeeper.ZookeeperManager;
import com.wzf.greattruth.zookeeper.dto.CreateNodeParam;
import com.wzf.greattruth.zookeeper.helper.ZookeeperHelper;
import com.wzf.greattruth.zookeeper.lock.DistributedLock;
import com.wzf.greattruth.zookeeper.node.Node;

public class DistributedQueue {

    private static final String QUEUE_PREFIX = "que-" ;
    
    private ZooKeeper zookeeper;
    /**
     * 分布式队列的目录
     */
    private String queueDir;

    private ZookeeperManager zookeeperManager;

    /**
     * 创建一个分布式队列
     * 
     * @param zookeeper
     * @param queueDir
     */
    public DistributedQueue(ZooKeeper zookeeper, String queueDir) {
        super();
        this.zookeeper = zookeeper;
        this.queueDir = queueDir;
        zookeeperManager = ZookeeperManager.getInstance();
    }

    /**
     * 向队列中添加一个数据
     * 
     * @param datas
     * @return
     */
    public String offer(byte[] datas) {
        CreateNodeParam createNodeParam = new CreateNodeParam();
        createNodeParam.setDatas(datas);
        boolean createParentResult = zookeeperManager.createDirIfNotExist(zookeeper, queueDir);
        if (!createParentResult) {
            return null;
        }

        String path = ZookeeperHelper.getNodePath(queueDir, QUEUE_PREFIX) ;
        return zookeeperManager.createDir(zookeeper, path, datas, CreateMode.PERSISTENT_SEQUENTIAL);
    }

    /**
     * 从队列头取出数据
     */
    public byte[] peek() {
        byte[] datas = null;
        Node node = zookeeperManager.getMinNode(zookeeper, queueDir, QUEUE_PREFIX);
        datas = zookeeperManager.getDatas(zookeeper, node);
        return datas;
    }

    /**
     * 从队列头取出数据后从队列中删除此数据<br/>
     * 此方法会尝试将节点删除，但并不能保证删除成功，所以调用方注意做幂等
     * 
     */
    public byte[] take() {
        Node node = zookeeperManager.getMinNode(zookeeper, queueDir, QUEUE_PREFIX);
        byte[] datas = zookeeperManager.getDatas(zookeeper, node, true);
        return datas;
    }

    /**
     * 从队列头取出数据后从队列中删除此数据
     * <li>此方法会尝试将节点删除，但并不能保证删除成功，所以调用方注意做幂等
     * <li>此方法会上锁，如果client都通过此方法获取数据时，此方法放着在当前client执行完前，没有client能从zk获取数据
     * 
     * @param lockDir
     * @param lockKey
     * @return
     */
    public byte[] take(String lockDir, String lockKey) {
        DistributedLock locker = null;
        if (StringUtils.isNotBlank(lockDir) && StringUtils.isNotBlank(lockKey)) {
            locker = new DistributedLock(zookeeper);
        }
        byte[] datas = null;
        if (locker != null) {
            boolean lock = locker.lock(lockDir, lockKey);
            if (!lock) {
                return null;
            }
            try {
                datas = take();
            } finally {
                locker.unlock(lockDir, lockKey);
            }
        } else {
            datas = take();
        }
        return datas;
    }

}
